import numpy as np
import matplotlib.pyplot as plt
from matplotlib import cm
from mpl_toolkits.mplot3d import Axes3D
import matplotlib as mpl
import os

# 设置中文显示
plt.rcParams['font.sans-serif'] = ['SimHei']  # 用来正常显示中文标签
plt.rcParams['axes.unicode_minus'] = False    # 用来正常显示负号

# 创建复平面网格
x = np.linspace(-2, 2, 300)
y = np.linspace(-2, 2, 300)
X, Y = np.meshgrid(x, y)
Z = X + 1j*Y

# 计算函数值
W = Z**2
M = np.abs(W)  # 模长
Phase = np.angle(W)  # 相位 [-π, π]

# 创建三维图形
fig = plt.figure(figsize=(10, 8), dpi=100)  # 设置基础DPI
ax = fig.add_subplot(111, projection='3d')

# 设置标题
fig.suptitle('复变函数 $w = z^2$ 的三维模长曲面', fontsize=18)

# 归一化相位到[0,1]范围
norm_phase = (Phase + np.pi) / (2 * np.pi)
colors = cm.hsv(norm_phase)

# 绘制模长曲面，使用相位着色
surf = ax.plot_surface(X, Y, M, facecolors=colors, 
                      rstride=1, cstride=1, 
                      antialiased=True, alpha=0.95)

# 设置坐标轴标签
ax.set_xlabel('实部 Re(z)', labelpad=12)
ax.set_ylabel('虚部 Im(z)', labelpad=12)
ax.set_zlabel('模长 |w|', labelpad=12)

# 设置视角
ax.view_init(elev=30, azim=45)

# 添加相位颜色条
mappable = cm.ScalarMappable(cmap=cm.hsv)
mappable.set_array(Phase)
cbar = fig.colorbar(mappable, ax=ax, shrink=0.6, aspect=10, pad=0.1)
cbar.set_label('相位 (弧度)')

# 添加图例说明
ax.text2D(0.05, 0.95, "曲面高度表示模长 |w| = |z|²\n颜色表示相位 arg(w)", 
          transform=ax.transAxes, fontsize=12,
          bbox=dict(facecolor='white', alpha=0.8))

# 添加坐标轴箭头
ax.quiver(0, 0, 0, 2.5, 0, 0, color='k', arrow_length_ratio=0.05, linewidth=1)  # x轴
ax.quiver(0, 0, 0, 0, 2.5, 0, color='k', arrow_length_ratio=0.05, linewidth=1)  # y轴
ax.quiver(0, 0, 0, 0, 0, 5, color='k', arrow_length_ratio=0.05, linewidth=1)    # z轴

# 在原点添加特殊标记
ax.scatter(0, 0, 0, c='black', s=50, marker='o', edgecolors='white')

# 设置坐标轴范围
ax.set_xlim(-2, 2)
ax.set_ylim(-2, 2)
ax.set_zlim(0, 8)

# 添加网格线
ax.xaxis.pane.fill = False
ax.yaxis.pane.fill = False
ax.zaxis.pane.fill = False
ax.grid(True, linestyle='--', alpha=0.5)

# 调整布局
plt.tight_layout()
plt.subplots_adjust(top=0.95)

# 保存高分辨率图像
output_dir = "output_images"
if not os.path.exists(output_dir):
    os.makedirs(output_dir)

# 保存为300dpi的PNG文件
filename = os.path.join(output_dir, "复变函数_z2_模长曲面_相位着色.png")
plt.savefig(filename, dpi=300, bbox_inches='tight', pad_inches=0.2)
print(f"图像已保存为: {filename} (300dpi)")

# 显示图像
plt.show()